package com.swallow.auth.infrastructure.acl.oauth2.config.jwt;

import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.ImmutableJWKSet;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import com.swallow.auth.common.enums.ErrorCode;
import com.swallow.auth.common.exception.SysException;
import com.swallow.auth.infrastructure.acl.oauth2.constants.ClaimConstants;
import com.swallow.auth.infrastructure.mysql.po.AccountPO;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.ClassPathResource;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.NimbusJwtDecoder;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.config.annotation.web.configuration.OAuth2AuthorizationServerConfiguration;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.security.oauth2.server.resource.authentication.JwtAuthenticationConverter;
import org.springframework.security.oauth2.server.resource.authentication.JwtGrantedAuthoritiesConverter;
import org.springframework.security.rsa.crypto.KeyStoreKeyFactory;

import java.security.KeyPair;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;

/**
 * @author: yangjie.deng@resico.cn
 * @since: 2024-05-13 14:31:22
 * @version: v1.0.0
 * @describe:
 *
 * -- 生成公钥
 * keytool -genkeypair -keyalg RSA -dname "cn=18310740596@163.com,ou=swallow,o=燕子商城,l=重庆,st=重庆,c=CN" -alias swallowkey -validity 3650 -keystore swallow.jks -storepass Swallow@2024 -keypass Swallow@2024
 *
 *
 * -- 查看公私钥
 * keytool -list -rfc --keystore swallow.jks | openssl x509 -inform pem -pubkey
 */
@Configuration
public class JwtConfig {
    @Bean
    public OAuth2TokenCustomizer<JwtEncodingContext> oAuth2TokenCustomizer() {
        return context -> {
            // 检查登录用户信息是不是UserDetails，排除掉没有用户参与的流程
            if (context.getPrincipal().getPrincipal() instanceof UserDetails user) {

                if (user instanceof AccountPO account) {
                    // 在token 存储 用户编号、租户、权限等信息
                    JwtClaimsSet.Builder claims = context.getClaims();
                    claims.claim(ClaimConstants.AUTHORITY_FLAG, account.getAccountNo());
                    claims.claim(ClaimConstants.ACCOUNT_NO, account.getAccountNo());
                    claims.claim(ClaimConstants.ACCOUNT_NAME, account.getStaffName());
                    claims.claim(ClaimConstants.TENANT_NO, account.getTenantNo());
                } else {
                    throw new SysException(ErrorCode.SYSTEM_ERROR);
                }
            }
        };
    }

    /**
     * 自定义jwt解析器，设置解析出来的权限信息的前缀与在jwt中的key
     *
     * @return jwt解析器 JwtAuthenticationConverter
     */
    @Bean
    public JwtAuthenticationConverter jwtAuthenticationConverter() {
        JwtGrantedAuthoritiesConverter grantedAuthoritiesConverter = new JwtGrantedAuthoritiesConverter();
        // 设置解析权限信息的前缀，设置为空是去掉前缀
        grantedAuthoritiesConverter.setAuthorityPrefix("");
        // 设置权限信息在jwt claims中的key
        grantedAuthoritiesConverter.setAuthoritiesClaimName("authorities");

        JwtAuthenticationConverter jwtAuthenticationConverter = new JwtAuthenticationConverter();
        jwtAuthenticationConverter.setJwtGrantedAuthoritiesConverter(grantedAuthoritiesConverter);
        return jwtAuthenticationConverter;
    }

    /**
     * 配置jwk源，使用非对称加密，公开用于检索匹配指定选择器的JWK的方法
     *
     * @return JWKSource
     */
    @Bean
    public JWKSource<SecurityContext> jwkSource() {
        KeyPair keyPair = generateRsaKey();
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        RSAKey rsaKey = new RSAKey.Builder(publicKey)
                .privateKey(privateKey)
                .keyID("Swallow_Jwt_Key_001")
                .build();

        JWKSet jwkSet = new JWKSet(rsaKey);
        return new ImmutableJWKSet<>(jwkSet);
    }

    private static final String KEY_ALIAS = "swallowkey";
    private static final String KEY_PWD = "Swallow@2024";

    /**
     * 生成rsa密钥对，提供给jwk
     *
     * @return 密钥对
     */
    private static KeyPair generateRsaKey() {
        KeyPair keyPair;
        try {
            ClassPathResource classPathResource = new ClassPathResource("/jwt/swallow.jks");
            KeyStoreKeyFactory storeKeyFactory = new KeyStoreKeyFactory(classPathResource, KEY_PWD.toCharArray());
            keyPair = storeKeyFactory.getKeyPair(KEY_ALIAS, KEY_PWD.toCharArray());
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
        return keyPair;
    }

    /**
     * 配置jwt解析器
     *
     * @param jwkSource jwk源
     * @return JwtDecoder
     */
    @Bean
    public JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource, OAuth2AuthorizationService authorizationService) {
        JwtDecoder jwtDecoder = OAuth2AuthorizationServerConfiguration.jwtDecoder(jwkSource);
        NimbusJwtDecoder decoder = (NimbusJwtDecoder) jwtDecoder;
        // 注入自定义校验器，扩展 /oauth2/revocation,支持token失效
        decoder.setJwtValidator(new CustomerJwtValidator(authorizationService));
        return jwtDecoder;
    }
}
